import numpy as np

from ..sensor import RangeBearingSensor
from ..map import Map


class MappingExtendedKalmanFilter:
    def __init__(self, senor: RangeBearingSensor, map_: Map, P_est: np.ndarray, V_est: np.ndarray,
                 W_est: np.ndarray) -> None:
        super().__init__()

        self._sensor = senor
        self._map = map_

        self._P_est = P_est
        self._V_est = V_est
        self._W_est = W_est

    def step(self, xv, z_measure):
        for landmark_id, iz in z_measure:
            z_meas = iz[:2]

            if not self._map.is_in(landmark_id):
                g = self._sensor.G(xv, z_meas)

                self.add_P_est(xv, g, z_meas)

                self._map.add(landmark_id, g)

            lm_idx = self._map.get_index(landmark_id)
            g = self._map.get_feature(landmark_id)

            landmark_est = self._map.get_features().reshape(1, -1).flatten()

            z_pred = self._sensor.H(xv, g)
            hh = np.zeros((2, len(landmark_est)))
            hh[:, 2 * lm_idx: 2 * (lm_idx + 1)] = self._sensor.Hxi(xv, g)

            S = hh @ self._P_est @ hh.T + self._W_est
            K = self._P_est @ hh.T @ np.linalg.inv(S)

            y = self._sensor.handle_measure(z_meas - z_pred)
            landmark_est[:] = landmark_est + K @ y
            self._map.update(landmark_est.reshape(-1, 2))
            self._P_est[:, :] = self._P_est - K @ hh @ self._P_est

    def add_P_est(self, xv, xi, z):
        n = self._P_est.shape[0]
        Yz = np.zeros((n + 2, n + 2))
        Yz[:n, :n] = np.eye(n)
        Yz[-2:, -2:] = self._sensor.Gz(xv, z)
        P_1 = np.zeros((n + 2, n + 2))
        P_1[:n, :n] = self._P_est
        P_1[-2:, -2:] = self._W_est
        self._P_est = Yz @ P_1 @ Yz.T
